import usb.core
import usb.util
import numpy as np
import cv2
import threading
import time

class ElectronLowLevel:
    def __init__(self, vid=0x1001, pid=0x8023):
        # 当前所使用的接口号
        self.interface_number = None
        self.USB_VID = vid
        self.USB_PID = pid
        self.is_connected = False
        self.time_stamp = 0
        
        # USB endpoints
        self.EP0 = 0x00
        self.EP1_IN = 0x81 
        self.EP1_OUT = 0x01
        
        # Buffers
        self.ping_pong_write_index = 0
        self.usb_buffer_200 = bytearray(200)
        self.frame_buffer_tx = [bytearray(240 * 240 * 3) for _ in range(2)]
        self.extra_data_buffer_tx = [bytearray(32) for _ in range(2)]
        self.extra_data_buffer_rx = bytearray(32)
        
        self.device = None
        self._sync_thread = None
        self._stop_thread = False

    def connect(self):
        """连接设备"""
        try:
            # 查找设备
            self.device = usb.core.find(idVendor=self.USB_VID, idProduct=self.USB_PID)
            if self.device is None:
                return False

            # 设置默认配置
            self.device.set_configuration()
            cfg = self.device.get_active_configuration()

            # 自动扫描每个接口的 altsetting，找到包含 Bulk IN/OUT 的 InterfaceAltSetting
            ep_in = None
            ep_out = None
            sel_intf = None
            sel_alt = None
            for interface in cfg:
                # interface 代表某个接口的一个 altsetting
                for ep in interface:
                    if usb.util.endpoint_direction(ep.bEndpointAddress) == usb.util.ENDPOINT_IN \
                       and usb.util.endpoint_type(ep.bmAttributes) == usb.util.ENDPOINT_TYPE_BULK:
                        ep_in = ep
                    if usb.util.endpoint_direction(ep.bEndpointAddress) == usb.util.ENDPOINT_OUT \
                       and usb.util.endpoint_type(ep.bmAttributes) == usb.util.ENDPOINT_TYPE_BULK:
                        ep_out = ep
                if ep_in and ep_out:
                    sel_intf = interface.bInterfaceNumber
                    sel_alt = interface.bAlternateSetting
                    # 切换到对应的 altsetting
                    try:
                        self.device.set_interface_altsetting(interface=sel_intf, alternate_setting=sel_alt)
                        print(f"[Debug] 切换到接口 {sel_intf} 的 altsetting {sel_alt}")
                    except Exception as e:
                        print(f"[Warning] 切换 altsetting 失败: {e}")
                    # 卸载内核驱动并 Claim 接口
                    try:
                        if self.device.is_kernel_driver_active(sel_intf):
                            self.device.detach_kernel_driver(sel_intf)
                    except Exception:
                        pass
                    usb.util.claim_interface(self.device, sel_intf)

                    # 设置最终端点地址
                    self.ep_in = ep_in
                    self.ep_out = ep_out
                    self.EP1_IN = ep_in.bEndpointAddress
                    self.EP1_OUT = ep_out.bEndpointAddress
                    break
            if not (ep_in and ep_out):
                print(f"[Warning] 未找到包含 Bulk IN/OUT 的 altsetting，使用默认 EP1_IN=0x{self.EP1_IN:02X}, EP1_OUT=0x{self.EP1_OUT:02X}")

            self.is_connected = True
            self.time_stamp = 0
            return True

        except Exception as e:
            print(f"连接错误: {e}")
            return False

    def disconnect(self):
        """断开连接"""
        if self._sync_thread and self._sync_thread.is_alive():
            self._stop_thread = True
            self._sync_thread.join()
            
        if self.is_connected:
            usb.util.dispose_resources(self.device)
            self.is_connected = False
            return True
        return False

    def sync(self):
        """同步数据"""
        if not self.is_connected:
            return False
            
        if self._sync_thread and self._sync_thread.is_alive():
            self._stop_thread = True
            self._sync_thread.join()
            
        self._stop_thread = False
        self._sync_thread = threading.Thread(target=self._sync_task)
        self._sync_thread.start()
        
        self.time_stamp += 1
        return True

    def _sync_task(self):
        """同步任务"""
        frame_buffer_offset = 0
        index = self.ping_pong_write_index
        self.ping_pong_write_index = 1 if self.ping_pong_write_index == 0 else 0
        
        for p in range(4):
            if self._stop_thread:
                break
                
            # 接收额外数据
            self.receive_packet(self.extra_data_buffer_rx, 1, 32)
            
            # 发送图像数据
            self.transmit_packet(
                self.frame_buffer_tx[index][frame_buffer_offset:frame_buffer_offset + 43008],
                84, 512)
            frame_buffer_offset += 43008
            
            # 发送帧尾和额外数据
            tail_data = bytearray(224)
            tail_data[:192] = self.frame_buffer_tx[index][frame_buffer_offset:frame_buffer_offset + 192]
            tail_data[192:] = self.extra_data_buffer_tx[index]
            self.transmit_packet(tail_data, 1, 224)
            frame_buffer_offset += 192

    def set_image_src(self, image):
        """设置显示图像"""
        if isinstance(image, str):
            # 从文件加载图像
            img = cv2.imread(image)
        else:
            # 直接使用numpy数组
            img = image
            
        if img is None:
            return
            
        # 调整图像大小和格式
        img = cv2.resize(img, (240, 240))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        
        # 复制到缓冲区
        np.copyto(np.frombuffer(self.frame_buffer_tx[self.ping_pong_write_index], 
                               dtype=np.uint8).reshape(240, 240, 3), img)

    def set_joint_angles(self, j1, j2, j3, j4, j5, j6, enable=False):
        """设置关节角度"""
        angles = [j1, j2, j3, j4, j5, j6]
        buffer = self.extra_data_buffer_tx[self.ping_pong_write_index]
        
        # 设置使能标志
        buffer[0] = 1 if enable else 0
        
        # 写入角度数据
        for j, angle in enumerate(angles):
            angle_bytes = np.array([angle], dtype=np.float32).tobytes()
            buffer[j * 4 + 1:j * 4 + 5] = angle_bytes

    def get_joint_angles(self):
        """获取关节角度"""
        angles = []
        for j in range(6):
            angle_bytes = self.extra_data_buffer_rx[j * 4 + 1:j * 4 + 5]
            angle = np.frombuffer(angle_bytes, dtype=np.float32)[0]
            angles.append(angle)
        return angles

    def receive_packet(self, buffer, packet_count, packet_size):
        """接收数据包"""
        try:
            for _ in range(packet_count):
                data = self.device.read(self.EP1_IN, packet_size, timeout=100)
                buffer[:len(data)] = data
            return True
        except Exception as e:
            print(f"接收错误: {e}")
            return False

    def transmit_packet(self, buffer, packet_count, packet_size):
        """发送数据包"""
        try:
            for i in range(packet_count):
                offset = i * packet_size
                self.device.write(self.EP1_OUT, buffer[offset:offset + packet_size], timeout=100)
            return True
        except Exception as e:
            print(f"发送错误: {e}")
            return False